import os

import bs4
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain_chroma import Chroma
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_core.tools import create_retriever_tool
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.chat_message_histories import ChatMessageHistory
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent

os.environ["LANGCHAIN_TRACING_V2"] = "false"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ["LANGCHAIN_API_KEY"] = "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"
llm = ChatOpenAI(
    api_key="sk-CftUbVSsA61lwwgMz9xvt6znTunQZfgBP8ZCVLbQsKfXUR6k",
    model='deepseek-ai/DeepSeek-V3',
    base_url="https://www.henapi.top/v1"
)

loader = WebBaseLoader(
    web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
    bs_kwargs=dict(
        parse_only=bs4.SoupStrainer(class_=("post-content", "post-title", "post-header"))
    )
)
docs = loader.load()
# 创建文本分割器，设置每个chunk大小为1000字符，重叠200字符
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
# 将加载的文档分割成小块
splits = text_splitter.split_documents(docs)
# 使用OpenAI的嵌入模型将文档块转换为向量，并存储在Chroma向量数据库中
openAIEmbeddings = OpenAIEmbeddings(api_key="sk-CftUbVSsA61lwwgMz9xvt6znTunQZfgBP8ZCVLbQsKfXUR6k",
                                    model='text-embedding-3-small',
                                    base_url="https://www.henapi.top/v1")
vectorstore = Chroma.from_documents(documents=splits, embedding=openAIEmbeddings)
# 从向量数据库创建检索器，用于后续的相似性搜索
retriever = vectorstore.as_retriever()
system_prompt = (
    "You are an assistant for question-answering tasks. "
    "Use the following pieces of retrieved context to answer "
    "the question. If you don't know the answer, say that you "
    "don't know. Use three sentences maximum and keep the "
    "answer concise."
    "\n\n"
    "{context}"
)

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        ("human", "{input}"),
    ]
)

question_answer_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)
# response = rag_chain.invoke({"input": "What is Task Decomposition?"})
# print(response["answer"])

contextualize_q_system_prompt = (
    "Given a chat history and the latest user question "
    "which might reference context in the chat history, "
    "formulate a standalone question which can be understood "
    "without the chat history. Do NOT answer the question, "
    "just reformulate it if needed and otherwise return it as is."
)

contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)
history_aware_retriever = create_history_aware_retriever(
    llm, retriever, contextualize_q_prompt
)

qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)
question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

# chat_history = []
# question = "what is task decomposition?"
# ai_msg_1 = rag_chain.invoke({"input": question, "chat_history": chat_history})
# print(ai_msg_1["answer"])
# chat_history.extend(
#     [
#         HumanMessage(content=question),
#         AIMessage(content=ai_msg_1["answer"]),
#     ]
# )
# second_question = "What are common ways of doing it?"
# ai_msg_2 = rag_chain.invoke({"input": second_question, "chat_history": chat_history})
# print(ai_msg_2["answer"])

store = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


conversational_rag_chain = RunnableWithMessageHistory(
    rag_chain, get_session_history, input_messages_key="input", history_messages_key="chat_history",
    output_messages_key="answer",
)
# config={"configurable":{"session_id":"abc123"}}
# resp=conversational_rag_chain.invoke({"input": "What is Task Decomposition?"},config=config,)
# print(resp["answer"])
#
# resp=conversational_rag_chain.invoke(
#     {"input": "What are common ways of doing it?"},
#     config=config,
# )
#
# print(resp["answer"])


# for message in store["abc123"].messages:
#     if isinstance(message, AIMessage):
#         prefix = "AI"
#     else:
#         prefix = "User"
#
#     print(f"{prefix}: {message.content}\n")


tool = create_retriever_tool(
    retriever,
    "blog_post_retriever",
    "Searches and returns excerpts from the Autonomous Agents blog post.",
)
tools = [tool]

# resp=tool.invoke("task decomposition")
# print(resp)

config = {"configurable": {"thread_id": "abc123"}}
memory = MemorySaver()
agent_executor=create_react_agent(llm,tools,checkpointer=memory)
query="What is Task Decomposition?"
for s in agent_executor.stream(
        {"messages": [HumanMessage(content=query)]},config=config
):
    print(s)
    print("----")

query = "What are common ways of doing it?"

for s in agent_executor.stream(
    {"messages": [HumanMessage(content=query)]}, config=config
):
    print(s)
    print("----")